import pylab as plt
import os
import math
from params import *



dataset = 'gowalla'

degree = os.path.join(WORKDIR, 'gowalla', 'temporal', "results","%s_edge_degree_attachment.txt"%dataset)
age = os.path.join(WORKDIR, 'gowalla', 'temporal', "results","%s_edge_age_attachment_days.txt"%dataset)
distance = os.path.join(WORKDIR, 'gowalla', 'temporal', "results","%s_cpp_distance_attachment.txt"%dataset)

x = []
y = []
y2 = []
xmax = 1000
for line in open(degree):
    deg,new_k,weight_k,new_edges,new_weight = map(float, line.strip().split(','))
    if xmax and deg > xmax:
        continue
    if weight_k:
        x.append(deg)
        y.append(float(new_k)/weight_k)
        y2.append(float(new_k)/new_edges)

X1,X2 = 0,200
plt.figure()
plt.clf()
ax = plt.axes(FIG_AXES2)
plt.loglog(x,y,'ko',mfc="None")
#plt.loglog(x,y2,'kx',mfc="None")
xfit = map(math.log,x[X1:X2])
yfit = map(math.log,y[X1:X2])
a,b = plt.polyfit(xfit,yfit,1)
print "Fit: ",a,b
#a = 1.0
vfit = plt.polyval([a,b],map(math.log,x))
vfit = map(math.exp,vfit)
fit_leg = '$P_{deg}(k) \propto k^{%.2f}$'%a
fit_leg = '$k^{%.2f}$'%a
#fit_leg = '$k^%d$'%(int(a))
fx = plt.loglog(x,vfit,'k-',linewidth=2)
plt.legend([fx[0]], [fit_leg], loc='lower right')
plt.grid(True)
plt.xlabel('Node degree, $k$')
plt.ylabel('$P_{deg}(k)$')
#ymax = max(y)
#plt.ylim(ymax=ymax)
#if xmax:
#    x1,x2 = plt.xlim()
#    plt.xlim((x1,xmax))
plt.savefig('%s_degree_prob.pdf'%dataset)
plt.close()


#%plot_prob(age,'Node age (days), $t$', '$p_a(t)$',
#        'age_prob.pdf')#,xmax=350)#,fit=True)
x = []
y = []
#xmax = 350
for line in open(age):
    a,b = map(float, line.strip().split(','))
#    if xmax and a > xmax:
#        continue
    x.append(a)
    y.append(b)

plt.figure()
plt.clf()
ax = plt.axes(FIG_AXES2)
plt.semilogy(x,y,'ko',mfc="None")
plt.grid(True)
plt.xlabel('Node age (days), $a$')
plt.ylabel('$E(a)$')
plt.ylim(ymax=1.0)
plt.xlim(xmax=500)
#plt.xlim(xmax=70)
plt.savefig('%s_age_prob.pdf'%dataset)
plt.close()

#plot_prob(distance,'Link distance (km), $d$', '$p_g(d)$',
#        'distance_prob.pdf',xmax=1e4,log=True,fit=True,variable='d')

x = []
y = []
for line in open(distance):
    a,b = map(float, line.strip().split(','))
    if a >= 1e-2:
        if b:
            x.append(a)
            y.append(b)

plt.figure()
plt.clf()
ax = plt.axes(FIG_AXES2)
plt.loglog(x,y,'ko',mfc="None")
xfit = map(math.log10,x)
yfit = map(math.log10,y)
a,b = plt.polyfit(xfit,yfit,1)
print "Fit: ",a,b
a=-0.6
b=-3.2
#vfit = plt.polyval([a,b],map(math.log,x))
xfit = map(lambda x: 10**x,  [0, 4])
vfit = plt.polyval([a,b], [0, 4])
vfit = map(lambda x: 10**x, vfit)

print xfit, vfit, 
#fx = plt.loglog(xfit,vfit,'k--',linewidth=1)
#fit_leg = '$P_{geo}(d) \propto d^{%.2f}$'%a
#fit_leg = '$d^{%.2f}$'%a
#plt.legend([fx[0]], [fit_leg], loc='lower left')
plt.grid(True)
plt.xlabel('Link length (km), $d$')
plt.ylabel('$P_{geo}(d)$')
#plt.ylim(ymax=1.0)
plt.xlim(xmin=1.0)
plt.savefig('%s_distance_prob.pdf'%dataset)
plt.close()
